load("//bazel:api.bzl", "modular_py_binary", "modular_py_library", "pkg_files", "requirement", "strip_prefix")

package(default_visibility = ["//max:consumers"])

modular_py_library(
    name = "benchmark_serving_lib",
    srcs = glob(
        ["*.py"],
        exclude = [
            "serve_replay.py",
            "benchmark_throughput.py",
        ],
    ),
    data = [
        ":benchmark_config_yaml_files",
    ],
    imports = ["../.."],
    deps = [
        "//max/python/max/benchmark/benchmark_shared",
        "//max/python/max/diagnostics",
        requirement("aiohttp"),
        requirement("datasets"),
        requirement("huggingface-hub"),
        requirement("jinja2"),
        requirement("msgspec"),
        requirement("numpy"),
        requirement("pillow"),
        requirement("psutil"),
        requirement("pyyaml"),
        requirement("tqdm"),
        requirement("transformers"),
    ],
)

modular_py_binary(
    name = "benchmark_serving",
    srcs = ["benchmark_serving.py"],
    imports = ["../.."],
    deps = [
        "//max/python/max/benchmark/benchmark_shared",
        "//max/python/max/diagnostics",
        requirement("aiohttp"),
        requirement("datasets"),
        requirement("huggingface-hub"),
        requirement("numpy"),
        requirement("jinja2"),
        requirement("tqdm"),
        requirement("transformers"),
    ],
)

modular_py_binary(
    name = "benchmark_throughput",
    srcs = ["benchmark_throughput.py"],
    data = [
        ":yaml_files",
    ],
    imports = ["../.."],
    deps = [
        "//max/python/max/benchmark/benchmark_shared",
        "//max/python/max/engine",
        "//max/python/max/entrypoints:pipelines",
        "//max/python/max/pipelines/architectures",
        "//max/python/max/serve",
        requirement("click"),
        requirement("msgspec"),
        requirement("pyarrow"),
        requirement("pyyaml"),
        requirement("transformers"),
    ],
)

modular_py_binary(
    name = "serve_replay",
    srcs = ["serve_replay.py"],
    visibility = [":__subpackages__"],
    deps = [
        "//max/python/max/entrypoints:pipelines",
        "//max/python/max/entrypoints:replay_recording",
        requirement("click"),
        requirement("httpx"),
    ],
)

# TODO: This is maybe redundant since we already have yaml_files?
pkg_files(
    name = "benchmark_config_yaml_files",
    srcs = glob(["configs/*.yaml"]),
)

pkg_files(
    name = "python_files",
    srcs = glob(["**/*.py"]) + ["//max/python/max/benchmark/benchmark_shared:python_files"],
    prefix = "benchmark",
    strip_prefix = strip_prefix.from_pkg(""),
)

pkg_files(
    name = "yaml_files",
    srcs = glob(
        ["**/*.yaml"],
        exclude = ["serving-benchmarking/**"],
    ),
    prefix = "benchmark",
)
